
#include "scatter_elements_tiling.h"
#include "register/op_def_registry.h"
#include "graph/utils/type_utils.h"
#include "tiling/platform/platform_ascendc.h"
#include <iostream>
#include <string>

namespace optiling {
const uint32_t BLOCK_SIZE = 32;
const uint32_t BUFFER_NUM = 2;
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{
    ScatterElementsTilingData tiling;
    uint64_t ubSize;
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
    auto coreNum = ascendcPlatform.GetCoreNum();
    
    int64_t axis = *context->GetAttrs()->GetInt(0);
    tiling.set_axis(axis);

    const char * reduce = context->GetAttrs()->GetStr(1);
    std::string s_reduce = "add";
    // 0表示加法，1表示乘法
    if(s_reduce.compare(reduce) == 0) {
        tiling.set_reduce(0);
        printf("add detected\n");
    } else {
        tiling.set_reduce(1);
        printf("mul detected\n");
    }
    

    // Based on the input length and the number of inputs, the number of bytes of the input data type is obtained
    uint32_t inputNum = context->GetInputShape(0)->GetStorageShape().GetShapeSize();
    uint32_t typeLength = 0;
    ge::TypeUtils::GetDataTypeLength(context->GetInputDesc(0)->GetDataType(), typeLength);
    uint32_t inputLength = inputNum * typeLength;
    uint32_t inputBytes = inputLength / inputNum;

    uint32_t indexLength = context->GetInputShape(1)->GetStorageShape().GetShapeSize();
    // indexLength要32B对齐
    uint32_t indexLengthAligned = ((indexLength * typeLength + BLOCK_SIZE - 1) / BLOCK_SIZE * BLOCK_SIZE) / typeLength;
    tiling.set_indexLength(indexLength);
    tiling.set_indexLengthAligned(indexLengthAligned);

     uint32_t updatesLength = context->GetInputShape(1)->GetStorageShape().GetShapeSize();
    // indexLength要32B对齐
    uint32_t updatesLengthAligned = ((updatesLength * typeLength + BLOCK_SIZE - 1) / BLOCK_SIZE * BLOCK_SIZE) / typeLength;
    tiling.set_updatesLength(updatesLength);
    tiling.set_updatesLengthAligned(updatesLengthAligned);

    auto shape = context->GetInputShape(0)->GetOriginShape();
    uint32_t unitCount = 1;
    uint32_t stepSize = 1;
    uint32_t totalLength = 1;
    uint32_t firstDimSelf = 1;
    uint32_t secondDimSelf = 1;
    uint32_t thirdDimSelf = 1;
    for(int i = 0; i < shape.GetDimNum(); i++) {
        if(i < axis) {
            unitCount *= shape.GetDim(i);
        } else if(i > axis) {
            stepSize *= shape.GetDim(i);
        } else {
            totalLength *= shape.GetDim(i);
        }

        if(i == shape.GetDimNum() - 1) {
            thirdDimSelf = shape.GetDim(i);
        }
        
        if(i == shape.GetDimNum() - 2) {
            secondDimSelf = shape.GetDim(i);    
        }

        if(i == shape.GetDimNum() - 3) {
            firstDimSelf = shape.GetDim(i);    
        }
    }

    uint32_t firstDimIndex= 1;
    uint32_t secondDimIndex = 1;
    uint32_t thirdDimIndex = 1;

    auto IndexShape = context->GetInputShape(1)->GetOriginShape();

    for(int i = 0; i < IndexShape.GetDimNum(); i++) {
        if(i == IndexShape.GetDimNum() - 1) {
            thirdDimIndex = IndexShape.GetDim(i);
        }
        
        if(i == IndexShape.GetDimNum() - 2) {
            secondDimIndex = IndexShape.GetDim(i);    
        }

        if(i == IndexShape.GetDimNum() - 3) {
            firstDimIndex = IndexShape.GetDim(i);    
        }
    }

    uint32_t firstDimUpdates= 1;
    uint32_t secondDimUpdates = 1;
    uint32_t thirdDimUpdates = 1;

    auto UpdatesShape = context->GetInputShape(2)->GetOriginShape();

    for(int i = 0; i < UpdatesShape.GetDimNum(); i++) {
        if(i == UpdatesShape.GetDimNum() - 1) {
            thirdDimUpdates = UpdatesShape.GetDim(i);
        }
        
        if(i == UpdatesShape.GetDimNum() - 2) {
            secondDimUpdates = UpdatesShape.GetDim(i);    
        }

        if(i == UpdatesShape.GetDimNum() - 3) {
            firstDimUpdates = UpdatesShape.GetDim(i);    
        }
    }
    
    tiling.set_firstDimSelf(firstDimSelf);
    tiling.set_secondDimSelf(secondDimSelf);
    tiling.set_thirdDimSelf(thirdDimSelf);
    tiling.set_firstDimIndex(firstDimIndex);
    tiling.set_secondDimIndex(secondDimIndex);
    tiling.set_thirdDimIndex(thirdDimIndex);
    tiling.set_firstDimUpdates(firstDimUpdates);
    tiling.set_secondDimUpdates(secondDimUpdates);
    tiling.set_thirdDimUpdates(thirdDimUpdates);
    tiling.set_dimSize(shape.GetDimNum());

    printf("firstDimSelf的长度是:%u\n", firstDimSelf);
    printf("secondDimSelf的长度是:%u\n", secondDimSelf);
    printf("thirdDimSelf的长度是:%u\n", thirdDimSelf);
    printf("dimSize的长度是:%u\n", shape.GetDimNum());
    printf("firstDimIndex的长度是:%u\n", firstDimIndex);
    printf("secondDimIndex:%u\n", secondDimIndex);
    printf("thirdDimIndex:%u\n", thirdDimIndex);
    printf("firstDimUpdates:%u\n", firstDimUpdates);
    printf("secondDimUpdates:%u\n", secondDimUpdates);
    printf("thirdDimUpdates:%u\n", thirdDimUpdates);

    

    uint32_t ubDataNumber = 0;
    auto dt = context->GetInputTensor(0)->GetDataType();
    if(dt == ge::DT_INT8 || dt == ge::DT_INT32 || dt == ge::DT_FLOAT16) {
        // 2 more TBuf for int types
        ubDataNumber = 6;
    } else {
        ubDataNumber = 4;
    }
    // The number of 32B data blocks that can be used for each data. DOUBLE BUFFER is already counted here
    uint32_t tileBlockNum = (ubSize / BLOCK_SIZE / BUFFER_NUM) / ubDataNumber;
    uint32_t tileDataNum = (tileBlockNum * BLOCK_SIZE) / inputBytes;

    // Input data for 32B alignment
    uint32_t inputLengthAlgin32 = (((inputLength + BLOCK_SIZE - 1) / BLOCK_SIZE) * BLOCK_SIZE);
    // There is at least 32B of data on each core, satisfying several settings for several cores. The maximum number of audits is the actual number of audits
    // coreNum = (coreNum <  inputLengthAlgin32 / BLOCK_SIZE) ? coreNum : inputLengthAlgin32 / BLOCK_SIZE;
    // coreNum = (coreNum >= 1) ? coreNum : 1;
    coreNum = 1;
    uint32_t everyCoreInputBlockNum = inputLengthAlgin32 / BLOCK_SIZE / coreNum;
    uint32_t tailBlockNum = (inputLengthAlgin32 / BLOCK_SIZE) % coreNum;
    
    // Small chunks are calculated and sliced several times using the number of data on each core
    uint32_t smallCoreDataNum = everyCoreInputBlockNum * BLOCK_SIZE / inputBytes;
    uint32_t smallTileNum = everyCoreInputBlockNum / tileBlockNum;
    uint32_t finalSmallTileNum = (everyCoreInputBlockNum % tileBlockNum) == 0 ? smallTileNum : smallTileNum + 1;
    // Tail block calculation for small chunks of data
    uint32_t smallTailDataNum = smallCoreDataNum - (tileDataNum * smallTileNum);
    smallTailDataNum = smallTailDataNum == 0 ? tileDataNum : smallTailDataNum;
    // 总共需要处理的数据
    uint32_t realSmallCoreDataNum = inputNum;
    
    // The total length of a large block of data is 32B larger than that of a small block of data
    everyCoreInputBlockNum += 1;
    uint32_t bigCoreDataNum = everyCoreInputBlockNum * BLOCK_SIZE / inputBytes;
    uint32_t bigTileNum = everyCoreInputBlockNum / tileBlockNum;
    uint32_t finalBigTileNum = (everyCoreInputBlockNum % tileBlockNum) == 0 ? bigTileNum : bigTileNum + 1;
    uint32_t bigTailDataNum = bigCoreDataNum - tileDataNum * bigTileNum;
    bigTailDataNum = bigTailDataNum == 0 ? tileDataNum : bigTailDataNum;
    
    tiling.set_smallCoreDataNum(realSmallCoreDataNum);
    tiling.set_bigCoreDataNum(bigCoreDataNum);
    tiling.set_tileDataNum(tileDataNum);
    tiling.set_smallTailDataNum(smallTailDataNum);
    tiling.set_bigTailDataNum(bigTailDataNum);
    tiling.set_finalSmallTileNum(finalSmallTileNum);
    tiling.set_finalBigTileNum(finalBigTileNum);
    tiling.set_tailBlockNum(tailBlockNum);
    tiling.set_unitCount(unitCount);
    tiling.set_totalLength(totalLength);
    tiling.set_stepSize(stepSize);


    uint32_t indexTileNum = indexLengthAligned / tileDataNum;
    indexTileNum = indexLengthAligned % tileDataNum ? indexTileNum + 1 : indexTileNum;
    uint32_t indexTailDataNum = indexLengthAligned % tileDataNum;
    indexTailDataNum = indexTailDataNum == 0 ? tileDataNum : indexTailDataNum;

    uint32_t updatesTileNum = updatesLengthAligned / tileDataNum;
    updatesTileNum = updatesLengthAligned % tileDataNum ? updatesTileNum + 1 : updatesTileNum;
    uint32_t updatesTailDataNum = updatesLengthAligned % tileDataNum;
    updatesTailDataNum = updatesTailDataNum == 0 ? tileDataNum : updatesTailDataNum;

    tiling.set_indexTileNum(indexTileNum);
    tiling.set_indexTailDataNum(indexTailDataNum);
    tiling.set_updatesTileNum(updatesTileNum);
    tiling.set_updatesTailDataNum(updatesTailDataNum);

    printf("indexTileNum:%u\n", indexTileNum);
    printf("indexTailDataNum:%u\n", indexTailDataNum);
    printf("updatesTileNum:%u\n", updatesTileNum);
    printf("updatesTailDataNum:%u\n", updatesTailDataNum);

    printf("inputNum: %u\n", inputNum);
    printf("finalSmallTileNum: %u\n", finalSmallTileNum);
    printf("smallTailDataNum: %u\n", smallTailDataNum);
    printf("realSmallCoreDataNum: %u\n", realSmallCoreDataNum);
    printf("tileDataNum: %u\n", tileDataNum);
    printf("index数组的实际长度: %u\n", indexLength);
    printf("index数组的对齐长度: %u\n", indexLengthAligned);
    printf("updates 数组的对齐长度: %u\n", updatesLength);
    printf("updates 数组的对齐长度: %u\n", updatesLengthAligned);
    printf("stepSize的长度是:%u\n", stepSize);
    printf("totalLength的长度是:%u\n", totalLength);
    
    context->SetBlockDim(coreNum);
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    size_t *currentWorkspace = context->GetWorkspaceSizes(1);
    currentWorkspace[0] = 0;
    return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class ScatterElements : public OpDef {
public:
    explicit ScatterElements(const char* name) : OpDef(name)
    {
        this->Input("var")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT32, ge::DT_UINT8})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("indices")
            .ParamType(REQUIRED)
            .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("updates")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT32, ge::DT_UINT8})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Attr("axis").Int();
        this->Attr("reduce").AttrType(OPTIONAL).String("None");

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend910");
	this->AICore().AddConfig("ascend310b");

    }
};

OP_ADD(ScatterElements);
}

